""" This file contains Visualizer class based on Facebook's visdom.

Returns:
    Visualizer(): Visualizer class to display plots and images
"""
import datetime
##
import os
import time
import numpy as np
import torchvision.utils as vutils
import matplotlib.pyplot as plt
import torch
##
class Visualizer():
    """ Visualizer wrapper based on Visdom.

    Returns:
        Visualizer: Class file.
    """
    # pylint: disable=too-many-instance-attributes
    # Reasonable.

    ##
    def __init__(self, opt, model=None, dataloader=None, device=None):
        plt.rcParams.update({
            'text.usetex': False,  # 禁用 LaTeX，彻底避免衬线字体依赖
            'axes.unicode_minus': False,  # 解决负号显示问题
            # 仅指定核心字体列表，覆盖 sans-serif 和 serif（用系统已安装字体）
            'font.sans-serif': ['SimSun', 'Microsoft YaHei', 'Arial', 'DejaVu Sans'],
            'font.serif': ['SimSun', 'Microsoft YaHei', 'Arial', 'DejaVu Sans'],
        })
        # self.opt = opt
        self.display_id = opt.display_id
        self.win_size = 256
        self.name = opt.name
        self.opt = opt
        if self.opt.display:
            import visdom
            self.vis = visdom.Visdom(server=opt.display_server, port=opt.display_port)

        # --
        # Dictionaries for plotting data and results.
        self.plot_data = None
        self.plot_res = None
        self.dataloader = dataloader
        self.model = model
        self.device = device
        # --
        # Path to train and test directories.
        self.img_dir = os.path.join(opt.outf, opt.name, 'train', 'images')
        self.tst_img_dir = os.path.join(opt.outf, opt.name, 'test', 'images')
        # 新增：曲线保存目录
        self.plot_dir = os.path.join(opt.outf, opt.name, 'train', 'plots')
        if not os.path.exists(self.img_dir):
            os.makedirs(self.img_dir)
        if not os.path.exists(self.tst_img_dir):
            os.makedirs(self.tst_img_dir)
        if not os.path.exists(self.plot_dir):
            os.makedirs(self.plot_dir)
        # --
        # Log file.
        self.log_name = os.path.join(opt.outf, opt.name, 'loss_log.txt')
        with open(self.log_name, "a") as log_file:
            now = time.strftime("%c")
            log_file.write('================ Training Loss (%s) ================\n' % now)

    ##
    @staticmethod
    def normalize(inp):
        """Normalize the tensor

        Args:
            inp ([FloatTensor]): Input tensor

        Returns:
            [FloatTensor]: Normalized tensor.
        """
        return (inp - inp.min()) / (inp.max() - inp.min() + 1e-5)

    ##
    def plot_current_errors(self, epoch, counter_ratio, errors):
        """Plot current errros.

        Args:
            epoch (int): Current epoch
            counter_ratio (float): Ratio to plot the range between two epoch.
            errors (OrderedDict): Error for the current epoch.
        """

        if not hasattr(self, 'plot_data') or self.plot_data is None:
            self.plot_data = {'X': [], 'Y': [], 'legend': list(errors.keys())}
        self.plot_data['X'].append(epoch + counter_ratio)
        self.plot_data['Y'].append([errors[k] for k in self.plot_data['legend']])

        # 1. 原有Visdom实时可视化（保留）
        if self.opt.display and hasattr(self, 'vis'):
            self.vis.line(
                X=np.stack([np.array(self.plot_data['X'])] * len(self.plot_data['legend']), 1),
                Y=np.array(self.plot_data['Y']),
                opts={
                    'title': self.name + ' loss over time',
                    'legend': self.plot_data['legend'],
                    'xlabel': 'Epoch',
                    'ylabel': 'Loss'
                },
                win=4
            )

        # 2. 新增：保存损失曲线到本地（每轮更新）
        self.save_loss_curve()

    ##
    def save_loss_curve(self):
        """保存损失曲线到本地目录"""
        plt.figure(figsize=(10, 6))
        # 转换数据格式
        X = np.array(self.plot_data['X'])
        Y = np.array(self.plot_data['Y'])
        legends = self.plot_data['legend']

        # 绘制每条损失曲线
        colors = ['blue', 'red', 'green', 'orange', 'purple', 'brown']  # 支持最多6条曲线，可扩展
        for i, legend in enumerate(legends):
            color = colors[i % len(colors)]  # 循环使用颜色
            plt.plot(X, Y[:, i], label=legend, color=color, linewidth=2)

        # 设置图表样式
        plt.xlabel('Epoch', fontsize=12)
        plt.ylabel('Loss Value', fontsize=12)
        plt.title(f'{self.name} Training Loss Curve', fontsize=14, fontweight='bold')
        plt.legend(fontsize=10)
        plt.grid(True, alpha=0.3)
        plt.tight_layout()

        # 保存图片（覆盖式保存，始终保留最新曲线）
        save_path = os.path.join(self.plot_dir, 'loss_curve.png')
        plt.savefig(save_path, dpi=150, bbox_inches='tight')
        plt.close()  # 关闭图表，避免内存占用

    ##
    def plot_performance(self, epoch, counter_ratio, performance):
        """ Plot performance

        Args:
            epoch (int): Current epoch
            counter_ratio (float): Ratio to plot the range between two epoch.
            performance (OrderedDict): Performance for the current epoch.
        """
        if not hasattr(self, 'plot_res') or self.plot_res is None:
            self.plot_res = {'X': [], 'Y': [], 'legend': list(performance.keys())}
        self.plot_res['X'].append(epoch + counter_ratio)
        self.plot_res['Y'].append([performance[k] for k in self.plot_res['legend']])

        # 1. 原有Visdom实时可视化（保留）
        if self.opt.display and hasattr(self, 'vis'):
            self.vis.line(
                X=np.stack([np.array(self.plot_res['X'])] * len(self.plot_res['legend']), 1),
                Y=np.array(self.plot_res['Y']),
                opts={
                    'title': self.name + 'Performance Metrics',
                    'legend': self.plot_res['legend'],
                    'xlabel': 'Epoch',
                    'ylabel': 'Stats'
                },
                win=5
            )

        # 2. 新增：保存性能指标曲线到本地（每轮更新）
        self.save_performance_curve()

    ##
    def save_performance_curve(self):
        """保存性能指标曲线到本地目录"""
        plt.figure(figsize=(10, 6))
        # 转换数据格式
        X = np.array(self.plot_res['X'])
        Y = np.array(self.plot_res['Y'])
        legends = self.plot_res['legend']

        # 绘制每条性能曲线
        colors = ['darkblue', 'darkred', 'darkgreen', 'darkorange']
        for i, legend in enumerate(legends):
            color = colors[i % len(colors)]
            plt.plot(X, Y[:, i], label=legend, color=color, linewidth=2, marker='o', markersize=4)

        # 设置图表样式
        plt.xlabel('Epoch', fontsize=12)
        plt.ylabel('Metric Value', fontsize=12)
        plt.title(f'{self.name} Performance Metrics Curve', fontsize=14, fontweight='bold')
        plt.legend(fontsize=10)
        plt.grid(True, alpha=0.3)
        plt.tight_layout()

        # 保存图片（覆盖式保存）
        save_path = os.path.join(self.plot_dir, 'performance_curve.png')
        plt.savefig(save_path, dpi=150, bbox_inches='tight')
        plt.close()

    ##
    def print_current_errors(self, epoch, errors):
        """ Print current errors.

        Args:
            epoch (int): Current epoch.
            errors (OrderedDict): Error for the current epoch.
            batch_i (int): Current batch
            batch_n (int): Total Number of batches.
        """
        # message = '   [%d/%d] ' % (epoch, self.opt.niter)
        message = '   Loss: [%d/%d] ' % (epoch, self.opt.niter)
        for key, val in errors.items():
            message += '%s: %.3f ' % (key, val)

        print(message)
        with open(self.log_name, "a") as log_file:
            log_file.write('%s\n' % message)

    ##
    def print_current_performance(self, performance, best):
        """ Print current performance results.

        Args:
            performance ([OrderedDict]): Performance of the model
            best ([int]): Best performance.
        """
        message = '   '
        for key, val in performance.items():
            message += '%s: %.3f ' % (key, val)
        message += 'max ' + self.opt.metric + ': %.3f' % best

        print(message)
        with open(self.log_name, "a") as log_file:
            log_file.write('%s\n' % message)

    def display_current_images(self, reals, fakes, fixed):
        """ Display current images.

        Args:
            epoch (int): Current epoch
            counter_ratio (float): Ratio to plot the range between two epoch.
            reals ([FloatTensor]): Real Image
            fakes ([FloatTensor]): Fake Image
            fixed ([FloatTensor]): Fixed Fake Image
        """
        reals = self.normalize(reals.cpu().numpy())
        fakes = self.normalize(fakes.cpu().numpy())
        fixed = self.normalize(fixed.cpu().numpy())

        if self.opt.display and hasattr(self, 'vis'):
            self.vis.images(reals, win=1, opts={'title': 'Reals'})
            self.vis.images(fakes, win=2, opts={'title': 'Fakes'})
            self.vis.images(fixed, win=3, opts={'title': 'Fixed'})

    def save_current_images(self, epoch, reals, fakes, fixed):
        """ Save images for epoch i.

        Args:
            epoch ([int])        : Current epoch
            reals ([FloatTensor]): Real Image
            fakes ([FloatTensor]): Fake Image
            fixed ([FloatTensor]): Fixed Fake Image
        """
        vutils.save_image(reals, '%s/reals.png' % self.img_dir, normalize=True)
        vutils.save_image(fakes, '%s/fakes.png' % self.img_dir, normalize=True)
        vutils.save_image(fixed, '%s/fixed_fakes_%03d.png' %(self.img_dir, epoch+1), normalize=True)

    def plot_loss_curve(self, loss_history, performance_history):
        """生成并保存各类loss曲线和汇总图"""
        plt.style.use('seaborn-v0_8-darkgrid')
        loss_plot_dir = os.path.join(self.opt.outf, self.opt.name, 'train', 'loss_plots')
        os.makedirs(loss_plot_dir, exist_ok=True)

        # 1. 绘制主损失曲线（Generator和Discriminator总损失）
        plt.figure(figsize=(10, 6))
        plt.plot(loss_history['err_g'], label='Generator Loss (err_g)', linewidth=2)
        plt.plot(loss_history['err_d'], label='Discriminator Loss (err_d)', linewidth=2)
        if self.opt.use_semi and 'err_cls' in loss_history:
            plt.plot(loss_history['err_cls'], label='Classification Loss (err_cls)', linewidth=2)
        plt.title('Main Losses Over Time', fontsize=14)
        plt.xlabel('Epoch', fontsize=12)
        plt.ylabel('Loss', fontsize=12)
        plt.legend(fontsize=10)
        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        plt.savefig(os.path.join(loss_plot_dir, 'main_losses.png'), dpi=300, bbox_inches='tight')
        plt.close()

        # 2. 绘制Generator分量损失曲线
        plt.figure(figsize=(10, 6))
        plt.plot(loss_history['err_g_adv'], label='Adversarial Loss (err_g_adv)', linewidth=2)
        plt.plot(loss_history['err_g_con'], label='Reconstruction Loss (err_g_con)', linewidth=2)
        plt.plot(loss_history['err_g_enc'], label='Encoder Loss (err_g_enc)', linewidth=2)
        plt.title('Generator Component Losses', fontsize=14)
        plt.xlabel('Epoch', fontsize=12)
        plt.ylabel('Loss', fontsize=12)
        plt.legend(fontsize=10)
        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        plt.savefig(os.path.join(loss_plot_dir, 'generator_components.png'), dpi=300, bbox_inches='tight')
        plt.close()

        # 3. 绘制高低频损失曲线（如果有）
        if len(loss_history['err_g_con_low']) > 0 and len(self.loss_history['err_g_con_high']) > 0:
            plt.figure(figsize=(10, 6))
            plt.plot(loss_history['err_g_con_low'], label='Low Frequency Loss', linewidth=2)
            plt.plot(loss_history['err_g_con_high'], label='High Frequency Loss', linewidth=2)
            plt.title('Frequency Component Losses', fontsize=14)
            plt.xlabel('Epoch', fontsize=12)
            plt.ylabel('Loss', fontsize=12)
            plt.legend(fontsize=10)
            plt.grid(True, alpha=0.3)
            plt.tight_layout()
            plt.savefig(os.path.join(loss_plot_dir, 'frequency_losses.png'), dpi=300, bbox_inches='tight')
            plt.close()

        # 4. 绘制增强损失曲线（如果有）
        enhanced_losses = []
        if 'err_contrast' in loss_history and len(loss_history['err_contrast']) > 0:
            enhanced_losses.append('err_contrast')
        if 'err_push' in loss_history and len(loss_history['err_push']) > 0:
            enhanced_losses.append('err_push')
        if 'err_latent_reg' in loss_history and len(loss_history['err_latent_reg']) > 0:
            enhanced_losses.append('err_latent_reg')

        if enhanced_losses:
            plt.figure(figsize=(10, 6))
            for loss_name in enhanced_losses:
                loss_label = loss_name.replace('err_', '').replace('_', ' ').title() + ' Loss'
                plt.plot(loss_history[loss_name], label=loss_label, linewidth=2)
            plt.title('Enhanced Loss Components', fontsize=14)
            plt.xlabel('Epoch', fontsize=12)
            plt.ylabel('Loss', fontsize=12)
            plt.legend(fontsize=10)
            plt.grid(True, alpha=0.3)
            plt.tight_layout()
            plt.savefig(os.path.join(loss_plot_dir, 'enhanced_losses.png'), dpi=300, bbox_inches='tight')
            plt.close()

        # 5. 绘制性能曲线
        if len(performance_history['auc']) > 0:
            plt.figure(figsize=(10, 6))
            plt.plot(performance_history['auc'], label=f'Test {self.opt.metric.upper()}', linewidth=2,
                     color='green')
            plt.axhline(y=performance_history['best_auc'], color='red', linestyle='--',
                        label=f'Best {self.opt.metric.upper()}: {performance_history["best_auc"]:.4f}')
            plt.title(f'Model Performance ({self.opt.metric.upper()}) Over Time', fontsize=14)
            plt.xlabel('Epoch', fontsize=12)
            plt.ylabel(self.opt.metric.upper(), fontsize=12)
            plt.legend(fontsize=10)
            plt.grid(True, alpha=0.3)
            plt.tight_layout()
            plt.savefig(os.path.join(loss_plot_dir, 'performance_curve.png'), dpi=300, bbox_inches='tight')
            plt.close()

        # 6. 绘制汇总图 - 所有loss在一个图中（分多个子图）
        # 计算需要的子图数量
        n_plots = 2  # 主loss + generator分量
        if len(loss_history['err_g_con_low']) > 0:
            n_plots += 1
        if enhanced_losses:
            n_plots += 1
        if len(performance_history['auc']) > 0:
            n_plots += 1

        fig, axes = plt.subplots(n_plots, 1, figsize=(12, 5 * n_plots))
        if n_plots == 1:
            axes = [axes]

        # 子图1: 主loss
        axes[0].plot(loss_history['err_g'], label='Generator Loss', linewidth=2)
        axes[0].plot(loss_history['err_d'], label='Discriminator Loss', linewidth=2)
        if self.opt.use_semi and 'err_cls' in loss_history:
            axes[0].plot(loss_history['err_cls'], label='Classification Loss', linewidth=2)
        axes[0].set_title('Main Losses', fontsize=12)
        axes[0].set_xlabel('Epoch')
        axes[0].set_ylabel('Loss')
        axes[0].legend()
        axes[0].grid(True, alpha=0.3)

        # 子图2: Generator分量loss
        axes[1].plot(loss_history['err_g_adv'], label='Adversarial Loss', linewidth=2)
        axes[1].plot(loss_history['err_g_con'], label='Reconstruction Loss', linewidth=2)
        axes[1].plot(loss_history['err_g_enc'], label='Encoder Loss', linewidth=2)
        axes[1].set_title('Generator Component Losses', fontsize=12)
        axes[1].set_xlabel('Epoch')
        axes[1].set_ylabel('Loss')
        axes[1].legend()
        axes[1].grid(True, alpha=0.3)

        # 子图3: 高低频loss（如果有）
        idx = 2
        if len(loss_history['err_g_con_low']) > 0:
            axes[idx].plot(loss_history['err_g_con_low'], label='Low Frequency Loss', linewidth=2)
            axes[idx].plot(loss_history['err_g_con_high'], label='High Frequency Loss', linewidth=2)
            axes[idx].set_title('Frequency Component Losses', fontsize=12)
            axes[idx].set_xlabel('Epoch')
            axes[idx].set_ylabel('Loss')
            axes[idx].legend()
            axes[idx].grid(True, alpha=0.3)
            idx += 1

        # 子图4: 增强loss（如果有）
        if enhanced_losses:
            for loss_name in enhanced_losses:
                loss_label = loss_name.replace('err_', '').replace('_', ' ').title() + ' Loss'
                axes[idx].plot(loss_history[loss_name], label=loss_label, linewidth=2)
            axes[idx].set_title('Enhanced Loss Components', fontsize=12)
            axes[idx].set_xlabel('Epoch')
            axes[idx].set_ylabel('Loss')
            axes[idx].legend()
            axes[idx].grid(True, alpha=0.3)
            idx += 1

        # 子图5: 性能曲线（如果有）
        if len(performance_history['auc']) > 0:
            axes[idx].plot(performance_history['auc'], label=f'{self.opt.metric.upper()}', linewidth=2,
                           color='green')
            axes[idx].axhline(y=performance_history['best_auc'], color='red', linestyle='--',
                              label=f'Best: {performance_history["best_auc"]:.4f}')
            axes[idx].set_title(f'Model Performance ({self.opt.metric.upper()})', fontsize=12)
            axes[idx].set_xlabel('Epoch')
            axes[idx].set_ylabel(self.opt.metric.upper())
            axes[idx].legend()
            axes[idx].grid(True, alpha=0.3)

        plt.tight_layout()
        plt.savefig(os.path.join(loss_plot_dir, 'all_losses_summary.png'), dpi=300, bbox_inches='tight')
        plt.close()

        # 保存loss数据为numpy文件
        np.save(os.path.join(loss_plot_dir, 'loss_history.npy'), self.loss_history)
        # np.save(os.path.join(loss_plot_dir, 'freq_loss_history.npy'), self.freq_loss_history)
        np.save(os.path.join(loss_plot_dir, 'performance_history.npy'), self.performance_history)

        print(f"所有loss曲线已保存到: {loss_plot_dir}")


    def _tensor_to_numpy(self, tensor):
        """将张量转换为可视化的numpy数组"""
        np_img = tensor.squeeze().cpu().detach().permute(1, 2, 0).numpy()
        np_img = (np_img - np_img.min()) / (np_img.max() - np_img.min() + 1e-8)
        return np_img

    def visualize_freq_separation(self, x):
        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        """可视化高低频分离效果"""
        freq_vis_dir = os.path.join(self.opt.outf, self.opt.name, 'train', 'freq_visualization')
        os.makedirs(freq_vis_dir, exist_ok=True)

        # 获取一批数据
        data_iter = iter(self.dataloader['train'])
        data = next(data_iter)
        inputs, _ = data[0].to(self.device), data[1].to(self.device)

        # 获取freq-sep层
        if hasattr(self.model, 'netg') and hasattr(self.model.netg, 'encoder1') and 'freq-sep' in self.model.netg.encoder1.main._modules:
            freq_sep_layer = self.model.netg.encoder1.main._modules['freq-sep']

            # 对前5个样本进行高低频分离
            n_samples = min(5, inputs.size(0))
            fig, axes = plt.subplots(3, n_samples, figsize=(4 * n_samples, 12))

            for i in range(n_samples):
                img = inputs[i:i + 1]

                # 进行高低频分离
                low, high = freq_sep_layer.down_sep(img)

                # 上采样到原图大小
                low_up = torch.nn.functional.interpolate(low, size=img.shape[2:], mode='bilinear', align_corners=True)

                # 转换为可视化格式
                img_np = self._tensor_to_numpy(img)
                low_np = self._tensor_to_numpy(low_up)
                high_np = self._tensor_to_numpy(high)

                # 显示图像
                axes[0, i].imshow(img_np)
                axes[0, i].set_title(f'Original Image {i + 1}')
                axes[0, i].axis('off')

                axes[1, i].imshow(low_np)
                axes[1, i].set_title(f'Low Frequency {i + 1}')
                axes[1, i].axis('off')

                axes[2, i].imshow(high_np)
                axes[2, i].set_title(f'High Frequency {i + 1}')
                axes[2, i].axis('off')

            plt.tight_layout()
            plt.savefig(os.path.join(freq_vis_dir, f"freq_sep_epoch_{epoch}.png"),
                        dpi=300, bbox_inches='tight')
            plt.close()

            # 计算统计信息
            all_low = []
            all_high = []
            for data in self.dataloader['train']:
                inputs, _ = data[0].to(self.device), data[1].to(self.device)
                low, high = freq_sep_layer.down_sep(inputs)
                all_low.append(torch.mean(low).item())
                all_high.append(torch.mean(torch.abs(high)).item())
                if len(all_low) > 10:  # 只计算前10个batch
                    break

            # 保存统计信息
            stats_path = os.path.join(freq_vis_dir, 'frequency_stats.txt')
            with open(stats_path, 'a') as f:
                f.write(f"Epoch {epoch + 1}:\n")
                f.write(f"  Low frequency mean: {np.mean(all_low):.4f}\n")
                f.write(f"  High frequency mean: {np.mean(all_high):.4f}\n")
                f.write(f"  Low/High ratio: {np.mean(all_low) / np.mean(all_high):.4f}\n\n")

        else:
            print("未找到freq-sep层，无法可视化高低频分离效果")
